from querier.esquerier import ElasticSearchQuerier


class WeiboTrendAnalyzeQuerier(ElasticSearchQuerier):
    def __init__(self, es, index, doc_type):
        super(WeiboTrendAnalyzeQuerier, self).__init__(es, index, doc_type)
        # self.nlp_service = nlp_service

    def _build_query(self, args):
        user_id = args.get('user_id', None)
        from_date = args.get('from', None)
        to_date = args.get('to', None)
        if user_id is None:
            raise ValueError('"user_id" is needed.')

        query = self._genQuery(user_id, from_date, to_date)
        return query, {}, {'user_id': user_id, 'from': from_date, 'to': to_date}

    def _build_result(self, es_result, param):
        # total = es_result['hits']['total']
        agg = es_result['aggregations']
        data = extract_result(agg)
        return {
            'user_id': param['user_id'],
            'from': param['from'],
            'to': param['to'],
            'doc_counts': data['doc_counts'],
            'sum_likes': data['sum_likes'],
            'avg_likes': data['avg_likes'],
            'max_likes': data['max_likes'],

            'sum_retweets': data['sum_retweets'],
            'avg_retweets': data['avg_retweets'],
            'max_retweets': data['max_retweets'],

            'sum_comments': data['sum_comments'],
            'avg_comments': data['avg_comments'],
            'max_comments': data['max_comments'],

            'sum_engagement': data['sum_engagement'],
            'avg_engagement': data['avg_engagement'],
            'max_engagement': data['max_engagement'],

            'dates': data['dates']
        }

    def _genQuery(self, user_id, from_date, to_date):
        query = {
            "query": {
                "bool": {
                    "filter": [
                        {'term': {'user_id': user_id}},
                        {"range": {"publish_timestamp": {"from": from_date, "to": to_date}}}
                    ]
                }
            },
            "aggs": {
                "time_hist": {
                    "date_histogram": {
                        "field": "publish_timestamp",
                        "interval": "1d"
                    },
                    "aggs": {
                        "sum_likes": {"sum": {"field": "likes", "missing": 0}},
                        "max_likes": {"max": {"field": "likes", "missing": 0}},

                        "sum_retweets": {"sum": {"field": "retweets", "missing": 0}},
                        "max_retweets": {"max": {"field": "retweets", "missing": 0}},

                        "sum_comments": {"sum": {"field": "comments", "missing": 0}},
                        "max_comments": {"max": {"field": "comments", "missing": 0}},

                        "sum_engagement": {"sum": {"field": "sum_engagement", "missing": 0}},
                        "max_engagement": {"max": {"field": "sum_engagement", "missing": 0}},
                    }
                }
            },

            "size": 0
        }

        return query


def extract_result(agg):
    buckets = agg['time_hist']['buckets']
    doc_counts = []
    sum_likes = []
    max_likes = []
    avg_likes = []
    sum_retweets = []
    max_retweets = []
    avg_retweets = []
    sum_comments = []
    max_comments = []
    avg_comments = []

    sum_engagement = []
    max_engagement = []
    avg_engagement = []

    dates = []

    for b in buckets:
        doc_counts.append(b['doc_count'])
        sum_likes.append(b['sum_likes']['value'])
        avg_likes.append(0 if b['doc_count'] == 0 else b['sum_likes']['value'] / b['doc_count'])
        max_likes.append(0 if b['doc_count'] == 0 else b['max_likes']['value'])

        sum_retweets.append(b['sum_retweets']['value'])
        avg_retweets.append(0 if b['doc_count'] == 0 else b['sum_retweets']['value'] / b['doc_count'])
        max_retweets.append(0 if b['doc_count'] == 0 else b['max_retweets']['value'])

        sum_comments.append(b['sum_comments']['value'])
        avg_comments.append(0 if b['doc_count'] == 0 else b['sum_comments']['value'] / b['doc_count'])
        max_comments.append(0 if b['doc_count'] == 0 else b['max_comments']['value'])

        sum_engagement.append(b['sum_engagement']['value'])
        avg_engagement.append(0 if b['doc_count'] == 0 else b['sum_engagement']['value'] / b['doc_count'])
        max_engagement.append(0 if b['doc_count'] == 0 else b['max_engagement']['value'])

        dates.append(b['key_as_string'])

    return {
        'doc_counts': doc_counts,

        "sum_likes": sum_likes,
        "max_likes": max_likes,
        "avg_likes": avg_likes,

        "sum_retweets": sum_retweets,
        "max_retweets": max_retweets,
        "avg_retweets": avg_retweets,

        "sum_comments": sum_comments,
        "max_comments": max_comments,
        "avg_comments": avg_comments,

        'sum_engagement': sum_engagement,
        'max_engagement': max_engagement,
        'avg_engagement': avg_engagement,

        'dates': dates
    }


